import os
import numpy as np
from torchvision import transforms
from torch.utils.data.dataset import Dataset

def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='latin1')
    return dict
    
# def load_CIFAR10_train(ROOT):
#     X = []
#     Y = []
#     for b in range(1,6):
#         file_x = os.path.join(ROOT,"cifar-10-batches-py/data_batch_{}".format(b))
#         dict_x = unpickle(file_x)
#         X.append(dict_x["data"])
#         Y.append(dict_x["labels"])
#     Xtr = np.concatenate(X).reshape(50000, 3, 32, 32).transpose(0, 2, 3, 1).astype("float")
#     Ytr = np.concatenate(Y)
#     del X, Y
#     transform = transforms.ToTensor()
#     Xtr = transform(Xtr)

#     return (Xtr,Ytr)

# def load_CIFAR10_test(ROOT):
#     X = []
#     Y = []
#     file_x = os.path.join(ROOT,"cifar-10-batches-py/test_batch")
#     dict_x = unpickle(file_x)
#     X.append(dict_x["data"])
#     Y.append(dict_x["labels"])
#     Xte = np.concatenate(X).reshape(10000, 3, 32, 32).transpose(0, 2, 3, 1).astype("float")
#     Yte = np.concatenate(Y)
#     del X, Y
#     return Xte, Yte

class load_CIFAR10(Dataset):
    def __init__(self, path, train,transforms=None):
        X = []
        Y = []
        if train == True:
            for b in range(1,6):
                file_x = os.path.join(path,"cifar-10-batches-py/data_batch_{}".format(b))
                dict_x = unpickle(file_x)
                X.append(dict_x["data"])
                Y.append(dict_x["labels"])
        if train == False:
            file_x = os.path.join(path,"cifar-10-batches-py/test_batch")
            dict_x = unpickle(file_x)
            X.append(dict_x["data"])
            Y.append(dict_x["labels"])
            
        Xtr = np.concatenate(X).reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1).astype("float32")
        Ytr = np.concatenate(Y)

        self.data = Xtr
        self.labels = np.asarray(Ytr)
        self.transforms = transforms
 
    def __getitem__(self, index):
        single_image_label = self.labels[index]
        img_as_np = self.data[index]
        if self.transforms is not None:
            img_as_tensor = self.transforms(img_as_np)
        return (img_as_tensor, single_image_label)
 
    def __len__(self):
        return len(self.data)
